import torch

def print_state_dict_shapes(checkpoint_path):
    """
    打印权重文件中所有键及其对应的形状。
    :param checkpoint_path: 权重文件路径（.pth 或 .pt 文件）
    """
    try:
        # 加载权重文件
        state_dict = torch.load(checkpoint_path, map_location=torch.device('cpu'))
        
        # 如果权重文件是一个字典，直接提取 state_dict
        if isinstance(state_dict, dict) and 'state_dict' in state_dict:
            state_dict = state_dict['state_dict']
        
        # 打印键及其形状
        print(f"权重文件路径: {checkpoint_path}")
        print(f"总键数: {len(state_dict)}")
        print("键名及其形状：")
        for key, value in state_dict.items():
            print(f"  {key}: {value.shape}")
    
    except FileNotFoundError:
        print(f"错误：文件 {checkpoint_path} 未找到。")
    except Exception as e:
        print(f"加载权重文件时出错：{e}")

# 使用示例
if __name__ == "__main__":
    checkpoint_path = ""
    print_state_dict_shapes(checkpoint_path)